#!/bin/bash

# ==============================================================================
# Docker镜像导出脚本 - 完整版
# 包含所有docker-compose.yml和save_docker_images.sh中定义的镜像
# ==============================================================================

# 创建输出目录
OUTPUT_DIR="docker_images_backup"
mkdir -p $OUTPUT_DIR

# 定义输出文件名（包含时间戳）
OUTPUT_FILE="$OUTPUT_DIR/yuxi_know_docker_images_$(date +%Y%m%d_%H%M%S).tar"

echo "==============================================="
echo "开始导出 Yuxi-Know 项目所有 Docker 镜像"
echo "==============================================="
echo "输出文件: $OUTPUT_FILE"
echo ""

# 从 docker-compose.yml 提取的完整镜像列表
COMPOSE_IMAGES=(
    # 自定义构建镜像
    "yuxi-api:0.2.0.dev",
    "yuxi-web:0.2.0.dev",
    "mineru-sglang:latest",
    "paddlex:latest",

    # 数据库相关镜像
    "neo4j:5.26",
    "shangor/postgres-for-rag:v1.0",

    # Milvus 相关镜像
    "quay.io/coreos/etcd:v3.5.5",
    "minio/minio:RELEASE.2023-03-20T20-16-18Z",
    "milvusdb/milvus:v2.5.6",
)

# 从 save_docker_images.sh 提取的基础镜像列表
BASE_IMAGES=(
    "python:3.11-slim",
    "ghcr.io/astral-sh/uv:0.7.2",
    "node:20-alpine",
    "nginx:alpine",
    "neo4j:5.26",
    "quay.io/coreos/etcd:v3.5.5",
    "minio/minio:RELEASE.2023-03-20T20-16-18Z",
    "milvusdb/milvus:v2.5.6",
)

# 合并所有镜像并去重
ALL_IMAGES=(
    # 自定义构建镜像
    "yuxi-api:0.2.0.dev"
    "yuxi-web:0.2.0.dev"
    "mineru-sglang:latest"
    "paddlex:latest"

    # 数据库相关镜像
    "neo4j:5.26"
    "shangor/postgres-for-rag:v1.0"

    # Milvus 相关镜像
    "quay.io/coreos/etcd:v3.5.5"
    "minio/minio:RELEASE.2023-03-20T20-16-18Z"
    "milvusdb/milvus:v2.5.6"

    # 基础镜像（从Dockerfile分析得出）
    "python:3.11-slim"
    "ghcr.io/astral-sh/uv:0.7.2"
    "node:20-alpine"
    "nginx:alpine"

    # 构建依赖的基础镜像
    "lmsysorg/sglang:v0.4.9.post3-cu126"  # mineru.Dockerfile的基础镜像
    "ccr-2vdh3abv-pub.cnc.bj.baidubce.com/paddlex/paddlex:paddlex3.0.1-paddlepaddle3.0.0-gpu-cuda11.8-cudnn8.9-trt8.6"  # paddlex.Dockerfile的基础镜像
)

# 显示将要导出的镜像列表
echo "将要导出的镜像列表："
echo "-----------------------------------------------"
for i in "${!ALL_IMAGES[@]}"; do
    echo "$((i+1)). ${ALL_IMAGES[$i]}"
done
echo "总共 ${#ALL_IMAGES[@]} 个镜像"
echo ""

# 确认是否继续
read -p "确认导出这些镜像吗？(y/N): " -n 1 -r
echo ""
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "导出已取消"
    exit 0
fi

# 检查 Docker 是否运行
if ! docker info >/dev/null 2>&1; then
    echo "错误: Docker 未运行或权限不足"
    exit 1
fi

# 检查现有镜像并准备导出列表
echo "检查本地镜像..."
echo "-----------------------------------------------"
SUCCESS_IMAGES=()
MISSING_IMAGES=()

for IMAGE in "${ALL_IMAGES[@]}"; do
    # 使用更简单的检查方法：尝试在docker images输出中查找完全匹配的行
    if docker images --format "{{.Repository}}:{{.Tag}}" | grep -q "^${IMAGE}$"; then
        echo "✓ 找到镜像: $IMAGE"
        SUCCESS_IMAGES+=("$IMAGE")
    else
        echo "❌ 缺少镜像: $IMAGE"
        MISSING_IMAGES+=("$IMAGE")
    fi
done

echo ""
echo "镜像统计:"
echo "  ✓ 可用镜像: ${#SUCCESS_IMAGES[@]} 个"
echo "  ❌ 缺少镜像: ${#MISSING_IMAGES[@]} 个"

if [ ${#SUCCESS_IMAGES[@]} -eq 0 ]; then
    echo ""
    echo "错误: 没有找到任何可用的本地镜像"
    echo "建议: 请先运行 'docker compose build' 构建项目镜像"
    exit 1
fi

if [ ${#MISSING_IMAGES[@]} -gt 0 ]; then
    echo ""
    echo "缺少的镜像列表:"
    for IMAGE in "${MISSING_IMAGES[@]}"; do
        echo "  - $IMAGE"
    done
    echo ""
    echo "注意: 这些镜像将被跳过，如需包含请先拉取或构建"
fi

# 确认是否继续
echo ""
read -p "是否继续导出找到的 ${#SUCCESS_IMAGES[@]} 个镜像？(y/N): " -n 1 -r
echo ""
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "导出已取消"
    exit 0
fi

# 保存所有可用镜像到单个 tar 文件
echo ""
echo "开始保存镜像到 tar 文件..."
echo "-----------------------------------------------"

echo "正在保存 ${#SUCCESS_IMAGES[@]} 个镜像到 $OUTPUT_FILE..."
if docker save ${SUCCESS_IMAGES[@]} -o "$OUTPUT_FILE"; then
    echo "✓ 镜像保存成功"
else
    echo "❌ 镜像保存失败"
    exit 1
fi

# 计算文件大小
FILE_SIZE=$(du -h "$OUTPUT_FILE" | cut -f1)

# 生成镜像列表文件
LIST_FILE="$OUTPUT_DIR/image_list_$(date +%Y%m%d_%H%M%S).txt"
echo "生成镜像列表文件: $LIST_FILE"
echo "Yuxi-Know Docker 镜像导出列表" > "$LIST_FILE"
echo "导出时间: $(date)" >> "$LIST_FILE"
echo "镜像文件: $OUTPUT_FILE" >> "$LIST_FILE"
echo "文件大小: $FILE_SIZE" >> "$LIST_FILE"
echo "" >> "$LIST_FILE"
echo "包含的镜像:" >> "$LIST_FILE"
for IMAGE in "${SUCCESS_IMAGES[@]}"; do
    echo "  - $IMAGE" >> "$LIST_FILE"
done

# 生成加载脚本
LOAD_SCRIPT="$OUTPUT_DIR/load_images.sh"
echo "生成加载脚本: $LOAD_SCRIPT"
cat > "$LOAD_SCRIPT" << 'EOF'
#!/bin/bash

# Docker 镜像加载脚本
# 使用方法: ./load_images.sh [image_file.tar]

# 默认镜像文件
IMAGE_FILE=""
if [ $# -eq 1 ]; then
    IMAGE_FILE="$1"
else
    # 查找最新的 tar 文件
    IMAGE_FILE=$(ls -t docker_images_backup/*.tar 2>/dev/null | head -1)
fi

if [ -z "$IMAGE_FILE" ] || [ ! -f "$IMAGE_FILE" ]; then
    echo "错误: 找不到镜像文件"
    echo "使用方法: $0 <image_file.tar>"
    exit 1
fi

echo "正在从 $IMAGE_FILE 加载 Docker 镜像..."
echo ""

if docker load -i "$IMAGE_FILE"; then
    echo ""
    echo "✓ 镜像加载成功!"
    echo ""
    echo "加载的镜像列表:"
    docker images --format "table {{.Repository}}:{{.Tag}}\t{{.Size}}\t{{.CreatedAt}}" | head -n 20
else
    echo "❌ 镜像加载失败"
    exit 1
fi
EOF

chmod +x "$LOAD_SCRIPT"

# 最终输出
echo ""
echo "==============================================="
echo "导出完成！"
echo "==============================================="
echo "✓ 镜像文件: $OUTPUT_FILE"
echo "✓ 文件大小: $FILE_SIZE"
echo "✓ 镜像数量: ${#SUCCESS_IMAGES[@]}/${#ALL_IMAGES[@]}"
echo "✓ 镜像列表: $LIST_FILE"
echo "✓ 加载脚本: $LOAD_SCRIPT"
echo ""
echo "使用方法:"
echo "  1. 加载镜像: $LOAD_SCRIPT"
echo "  2. 或者直接: docker load -i $OUTPUT_FILE"
echo ""
echo "建议:"
echo "  - 将 $OUTPUT_DIR 目录复制到目标机器"
echo "  - 在目标机器上运行加载脚本"
echo "==============================================="

# 显示跳过的镜像（如果有）
SKIPPED_COUNT=$((${#ALL_IMAGES[@]} - ${#SUCCESS_IMAGES[@]}))
if [ $SKIPPED_COUNT -gt 0 ]; then
    echo ""
    echo "注意: 有 $SKIPPED_COUNT 个镜像未能拉取，未包含在导出文件中"
fi